import numpy as np
import pylab as plt
import matplotlib.pyplot as mp

### Methodes de choix du pas ###

def step_euler(y, t, h, f):
    return h*f(y,t) + y


def step_ptmilieu(y, t, h, f):
    return y + h*f( y + (h/2.) * f(y,t) , t+(h/2.) )


def step_heun(y,t,h,f):
    return y + (h/2.) * ( f(y,t) + f( y + h*f(y,t) , t+h ))


def step_RK4(y, t, h, f):
	k1 = f(y,t)
	k2 = f(y + (h/2.)*k1, t + h/2.)
	k3 = f(y + (h/2.)*k2, t + h/2.)
	k4 = f(y + h*k3 , t+h)
	return y + (h/6.)*(k1 + 2*k2 + 2*k3 + k4)


### Methodes du resolution ###


def meth_n_step(Yg,N,h,meth):
    t0 = Yg[0][0]
    y0 = Yg[0][1]
    sol=[y0]
    for i in range (0,N):
        ti = t0+i*h
        yi = sol[i]
        sol.append(meth(yi,ti,h,Yg[1]))
    return sol

def meth_epsilon(Yg, tf, eps, meth):
	t0 = Yg[0][0]
	y0 = Yg[0][1]
	f = Yg[1]
	i = 1
	h = np.linalg.norm(abs(tf-t0))
	tmp = np.linalg.norm(y0)
	sol = meth_n_step(Yg,1,h,meth)
	while(abs((sol[len(sol)-1]-tmp)) > eps):
		tmp = sol[len(sol)-1]
		i = i + 1
		h = abs(tf-t0)/i
		sol = meth_n_step(Yg,i,h,meth)
	return sol

### Champs de vecteurs ###

def champs_vecteurs(eq,ymin=0,ymax=2):
	f = eq[1]
	x = np.arange(-2., 2., 0.2)
	y = np.arange(0., 4., 0.2)
	X = []
	Y = []
	fx = []
	fy = []
	for i in range(len(x)):
		for j in range(len(y)):
			a = f(x[i],y[j])
			fy = fy + [a/(abs(a)*((np.sqrt((1./a)**2+1))))]
			fx = fx + [1./((np.sqrt(a**2+1)))]
			X = X + [x[i]]
			Y = Y + [y[j]]
	mp.clf()
	mp.ylim(ymin,ymax)
	mp.quiver(X,Y,fx,fy)
	x = np.arange(-2., 2., 0.1)
	mp.plot(x,meth_n_step(eq, 39, 0.1, step_RK4), label='Solution')
	mp.title('Champs des vecteurs tangents')
	mp.legend()
	mp.show()

if __name__ == "__main__":
	print "nothing to do here"
